#include <cstdint>
#include <cuda_runtime.h>
#include <cufft.h>
#include <npp.h>
#include <malloc.h>
#include <memory.h>
#include <memory>
#include "utils.h"
#include "ftlibcuda.h"
#include "cudafthandle.h"
#include "cudaFFTUtils.h"
#include "CudaTimer.h"

struct Deleter_CuFftPlan
{
	typedef cufftHandle pointer;

	void operator()(cufftHandle h)
	{
		if (h != NULL)
		{
			cufftDestroy(h);
		}
	}
};

struct Deleter_CudaMalloc
{
	typedef void* pointer;

	void operator()(void* p)
	{
		if (p != NULL)
		{
			cudaFree(p);
		}
	}
};

struct Deleter_AlignedMalloc
{
	typedef void* pointer;

	void operator()(void* p)
	{
		if (p != NULL)
		{
			_aligned_free(p);
		}
	}
};

/******************************************************************************/
// REAL-TO-COMPLEX (PACKED FFTW-Format)
/******************************************************************************/

CUDAFTPLANHANDLE cuFft_CreatePlan_2dRealToComplexPackedFftw(uint32_t width, uint32_t height, double normalization)
{
	CUDAFTPLANHANDLEDual handle;
	memset(&handle, 0, sizeof(handle));
	cufftHandle plan;
	cufftResult r = cufftPlan2d(&plan, width, height, CUFFT_R2C);
	ThrowIfCuFFTError(r, "cufftPlan2d");
	std::unique_ptr<cufftHandle, Deleter_CuFftPlan> cuFftPlan(plan);

	size_t srcSize = height * width * sizeof(float);
	void* devptrSrcData;
	cudaError_t cr = cudaMalloc(&devptrSrcData, srcSize);
	ThrowIfCudaError(cr, "cudaMalloc");
	std::unique_ptr<void*, Deleter_CudaMalloc> devPtrSrc(devptrSrcData);

	size_t destSize = height* /*width */(width / 2 + 1) * sizeof(float) * 2;
	void* devptrResultData;
	cr = cudaMalloc(&devptrResultData, destSize);
	ThrowIfCudaError(cr, "cudaMalloc");
	std::unique_ptr<void*, Deleter_CudaMalloc> devptrResult(devptrResultData);

	handle.internal.planType = CUDAFTPLANHANDLETYPE::_2dRealToComplexPackedFftw;
	handle.internal.ptrType1 = (intptr_t)devPtrSrc.release();
	handle.internal.ptrType2 = (intptr_t)devptrResult.release();
	handle.internal.cuFftplan = cuFftPlan.release();
	handle.internal.normalization = (float)normalization;

	return handle.external;
}

void cuFft_2d_RealToComplexPackedFftw(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfRealInputData2d& input, const FTLib::FourierTransfComplexOutput2d& output, FTLib::IPerformanceDataStore* perfData)
{
	CUDAFTPLANHANDLEDual& planX = (CUDAFTPLANHANDLEDual&)plan;
	ThrowIfPlanIsNotOfType(planX.internal, CUDAFTPLANHANDLETYPE::_2dRealToComplexPackedFftw);

	void* devptrSrcData = (void*)planX.internal.ptrType1;
	cudaError_t cr = cudaMemcpy2D(devptrSrcData, input.width*sizeof(float), input.ptrData, input.pitch, input.width*sizeof(float), input.height, cudaMemcpyHostToDevice);
	ThrowIfCudaError(cr, "cudaMemcpy2D");

	void* devptrDestComplex = (void*)planX.internal.ptrType2;
	cufftResult r = cufftExecR2C(planX.internal.cuFftplan, (cufftReal*)devptrSrcData, (cufftComplex*)devptrDestComplex);
	ThrowIfCuFFTError(r, "cufftExecR2C");

	if (planX.internal.normalization != 1)
	{
		NppiSize nppiSize = { (int)(output.width /*/ 2 + 1*/) * 2, (int)output.height };
		NppStatus status = nppiDivC_32f_C1IR(planX.internal.normalization, (Npp32f*)devptrDestComplex, (output.width /*/ 2 + 1*/) * sizeof(float) * 2, nppiSize);
		ThrowIfNppError(status, "nppiDivC_32f_C1IR");
	}

	cr = cudaMemcpy2D(output.ptrData, output.pitch, devptrDestComplex, output.width*sizeof(float) * 2, output.width * 2 * sizeof(float), output.height, cudaMemcpyDeviceToHost);
	ThrowIfCudaError(cr, "cudaMemcpy2D");
}

static void DeletePlan_PackedFftw(CUDAFTPLANHANDLEDual& planX)
{
	if (planX.internal.cuFftplan != NULL)
	{
		cufftResult r = cufftDestroy(planX.internal.cuFftplan);
	}

	if (planX.internal.ptrType1 != NULL)
	{
		cudaFree((void*)planX.internal.ptrType1);
	}

	if (planX.internal.ptrType2 != NULL)
	{
		cudaFree((void*)planX.internal.ptrType2);
	}
}

/******************************************************************************/
// REAL-TO-COMPLEX (Complex Format)
/******************************************************************************/

CUDAFTPLANHANDLE cuFft_CreatePlan_2dRealToComplex(uint32_t width, uint32_t height, double normalization)
{
	CUDAFTPLANHANDLEDual handle;
	memset(&handle, 0, sizeof(handle));
	cufftHandle plan;
	cufftResult r = cufftPlan2d(&plan, width, height, CUFFT_R2C);
	ThrowIfCuFFTError(r, "cufftPlan2d");
	std::unique_ptr<cufftHandle, Deleter_CuFftPlan> cuFftPlan(plan);
	uint32_t destHeight = height;
	uint32_t destWidth = width / 2 + 1;
	size_t destSize = destHeight * destWidth * sizeof(float) * 2;
	size_t srcSize = height * width * sizeof(float);

	void* devptrSrcData;
	cudaError_t cr = cudaMalloc(&devptrSrcData, srcSize);
	ThrowIfCudaError(cr, "cudaMalloc");
	std::unique_ptr<void*, Deleter_CudaMalloc> devPtrSrc(devptrSrcData);

	void* devptrDstData;
	cr = cudaMalloc(&devptrDstData, destSize);
	ThrowIfCudaError(cr, "cudaMalloc");
	std::unique_ptr<void*, Deleter_CudaMalloc> devPtrDst(devptrDstData);

	void* ptrHostTemp = _aligned_malloc(destSize, 32);
	ThrowIfNull(ptrHostTemp, "_aligned_malloc");
	std::unique_ptr<void*, Deleter_AlignedMalloc> hostPtrTemp(ptrHostTemp);

	handle.internal.cuFftplan = cuFftPlan.release();
	handle.internal.ptrType1 = (intptr_t)devPtrSrc.release();
	handle.internal.ptrType2 = (intptr_t)devPtrDst.release();
	handle.internal.ptrType3 = (intptr_t)hostPtrTemp.release();
	handle.internal.planType = CUDAFTPLANHANDLETYPE::_2dRealToComplex;
	handle.internal.normalization = (float)normalization;
	return handle.external;
}

template <class tPerf>
void cuFft_2d_RealToComplex(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfRealInputData2d& input, const FTLib::FourierTransfComplexOutput2d& output, tPerf& perf)
{
	CUDAFTPLANHANDLEDual& planX = (CUDAFTPLANHANDLEDual&)plan;
	ThrowIfPlanIsNotOfType(planX.internal, CUDAFTPLANHANDLETYPE::_2dRealToComplex);

	void* devptrSrcData = (void*)planX.internal.ptrType1;

	// we need 5 events at most -> host2Device, C2R, normalization, device2Host, finished
	perf.ReserveEvents(5);

	int host2DeviceEvent = perf.InsertCudaEvent();
	cudaError_t cr = cudaMemcpy2D(devptrSrcData, input.width*sizeof(float), input.ptrData, input.pitch, input.width*sizeof(float), input.height, cudaMemcpyHostToDevice);
	ThrowIfCudaError(cr, "cudaMemcpy2D");

	size_t destSize = output.height* (output.width / 2 + 1) * sizeof(float) * 2;

	void* devptrDestComplex = (void*)planX.internal.ptrType2;
	int cufftExecR2CEvent = perf.InsertCudaEvent();
	cufftResult r = cufftExecR2C(planX.internal.cuFftplan, (cufftReal*)devptrSrcData, (cufftComplex*)devptrDestComplex);
	/*ThrowIfCuFFTError(r, "cufftExecR2C");*/

	int normalizationEvent = -1;
	if (planX.internal.normalization != 1)
	{
		NppiSize nppiSize = { (int)(output.width / 2 + 1) * 2, (int)output.height };
		normalizationEvent = perf.InsertCudaEvent();
		NppStatus status = nppiDivC_32f_C1IR(planX.internal.normalization, (Npp32f*)devptrDestComplex, (output.width / 2 + 1) * sizeof(float) * 2, nppiSize);
		ThrowIfNppError(status, "nppiDivC_32f_C1IR");
	}

	float* fp = (float*)planX.internal.ptrType3;
	int device2HostEvent = perf.InsertCudaEvent();
	cr = cudaMemcpy(fp, devptrDestComplex, destSize, cudaMemcpyDeviceToHost);
	ThrowIfCudaError(cr, "cudaMemcpy2D");
	int finishedEvent = perf.InsertCudaEvent();

	// to my understanding, the cudaMemcpy-operation above will already synchronize operation,
	// so this statement is not required
	perf.SychnronizeToEvent(finishedEvent);

	perf.StartWallclock();
	ConvertFftwToComplex(input.width, input.height, fp, output.ptrData, output.pitch);
	perf.StopWallclock();
	perf.ReportDurationWallclock(L"Convert FFTW-Layout to Complex (on CPU)");

	perf.ReportDurationBetweenEvent(host2DeviceEvent, cufftExecR2CEvent, L"Copy Host->Device");
	if (normalizationEvent < 0)
	{
		// w/o normalization
		perf.ReportDurationBetweenEvent(cufftExecR2CEvent, device2HostEvent, L"cufftExecR2C");
		perf.ReportDurationBetweenEvent(device2HostEvent, finishedEvent, L"Copy Device->Host");
	}
	else
	{
		// with normalization
		perf.ReportDurationBetweenEvent(cufftExecR2CEvent, normalizationEvent, L"cufftExecR2C");
		perf.ReportDurationBetweenEvent(normalizationEvent, device2HostEvent, L"normalize");
		perf.ReportDurationBetweenEvent(device2HostEvent, finishedEvent, L"Copy Device->Host");
	}
}

void cuFft_2d_RealToComplex(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfRealInputData2d& input, const FTLib::FourierTransfComplexOutput2d& output, FTLib::IPerformanceDataStore* perfData)
{
	if (perfData == nullptr)
	{
		// the idea is to use a template with empty implementations, and hope that
		// the compiler's optimizer will throw out all related to benchmarking - although
		// I did not verify that really everything is thrown out...
		cuFft_2d_RealToComplex(plan, input, output, CNullPerfClass());
	}
	else
	{
		cuFft_2d_RealToComplex(plan, input, output, CPerfClass(0, perfData));
	}
}

static void DeletePlan_Complex(CUDAFTPLANHANDLEDual& planX)
{
	if (planX.internal.cuFftplan != NULL)
	{
		cufftResult r = cufftDestroy(planX.internal.cuFftplan);
	}

	if (planX.internal.ptrType1 != NULL)
	{
		cudaFree((void*)planX.internal.ptrType1);
	}

	if (planX.internal.ptrType2 != NULL)
	{
		cudaFree((void*)planX.internal.ptrType2);
	}

	if (planX.internal.ptrType3 != NULL)
	{
		_aligned_free((void*)planX.internal.ptrType3);
	}
}

/******************************************************************************/
// COMPLEX-TO-COMPLEX (Complex Format)
/******************************************************************************/

CUDAFTPLANHANDLE cuFft_CreatePlan_2dComplexToComplex(uint32_t width, uint32_t height, double normalization)
{
	CUDAFTPLANHANDLEDual handle;
	memset(&handle, 0, sizeof(handle));
	cufftHandle plan;
	cufftResult r = cufftPlan2d(&plan, width, height, CUFFT_C2C);
	ThrowIfCuFFTError(r, "cufftPlan2d");
	std::unique_ptr<cufftHandle, Deleter_CuFftPlan> cuFftPlan(plan);

	handle.internal.normalization = (float)normalization;
	size_t stride = width * sizeof(float) * 2;
	size_t size = height * stride;
	void* devptrData;
	cudaError_t cr = cudaMalloc(&devptrData, size);
	ThrowIfCudaError(cr, "cudaMalloc");
	handle.internal.ptrType1 = (intptr_t)devptrData;
	handle.internal.cuFftplan = cuFftPlan.release();
	handle.internal.planType = CUDAFTPLANHANDLETYPE::_2dComplexToComplex;
	return handle.external;
}

void cuFft_2d_ComplexToComplex(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfComplexInput2d& input, const FTLib::FourierTransfComplexOutput2d& output, bool forwardDirection, FTLib::IPerformanceDataStore* perfData)
{
	CUDAFTPLANHANDLEDual& planX = (CUDAFTPLANHANDLEDual&)plan;
	ThrowIfPlanIsNotOfType(planX.internal, CUDAFTPLANHANDLETYPE::_2dComplexToComplex);
	void* devptrData = (void*)planX.internal.ptrType1;
	auto stride = input.width * sizeof(float) * 2;
	cudaError_t cr = cudaMemcpy2D(devptrData, input.width*sizeof(float) * 2, input.ptrData, input.pitch, stride, input.height, cudaMemcpyHostToDevice);
	ThrowIfCudaError(cr, "cudaMemcpy2D");

	cufftResult r = cufftExecC2C(planX.internal.cuFftplan, (cufftComplex*)devptrData, (cufftComplex*)devptrData, (forwardDirection == true) ? CUFFT_FORWARD : CUFFT_INVERSE);
	ThrowIfCuFFTError(r, "cufftExecC2C");

	if (planX.internal.normalization != 1)
	{
		NppiSize nppiSize = { (int)input.width * 2, (int)input.height };
		NppStatus status = nppiDivC_32f_C1IR(planX.internal.normalization, (Npp32f*)devptrData, (int)stride, nppiSize);
		ThrowIfNppError(status, "nppiDivC_32f_C1IR");
	}

	cr = cudaMemcpy2D(output.ptrData, output.pitch, devptrData, stride, stride, output.height, cudaMemcpyDeviceToHost);
	ThrowIfCudaError(cr, "cudaMemcpy2D");
}

static void DeletePlan_Complex2Complex(CUDAFTPLANHANDLEDual& planX)
{
	cufftResult r = cufftDestroy(planX.internal.cuFftplan);
	cudaError_t cr = cudaFree((void*)planX.internal.ptrType1);
}

/******************************************************************************/
// HERMETIAN-TO-COMPLEX 
/******************************************************************************/

CUDAFTPLANHANDLE cuFft_CreatePlan_2dHermetianToReal(uint32_t width, uint32_t height, double normalization)
{
	CUDAFTPLANHANDLEDual handle;
	memset(&handle, 0, sizeof(handle));
	cufftHandle plan;
	cufftResult r = cufftPlan2d(&plan, width, height, CUFFT_C2R);
	ThrowIfCuFFTError(r, "cufftPlan2d");
	std::unique_ptr<cufftHandle, Deleter_CuFftPlan> cuFftPlan(plan);

	handle.internal.normalization = (float)normalization;
	size_t size = height * (width / 2 + 1) * sizeof(float) * 2;
	void* devptrData;
	cudaError_t cr = cudaMalloc(&devptrData, size);
	ThrowIfCudaError(cr, "cudaMalloc");
	handle.internal.ptrType1 = (intptr_t)devptrData;
	handle.internal.cuFftplan = cuFftPlan.release();
	handle.internal.planType = CUDAFTPLANHANDLETYPE::_2dHermetianToReal;
	return handle.external;
}

template <class tPerf>
void cuFft_2d_HermetianToReal(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfComplexInput2d& input, const FTLib::FourierTransfRealOutputData2d& output, tPerf& perf)
{
	CUDAFTPLANHANDLEDual& planX = (CUDAFTPLANHANDLEDual&)plan;
	ThrowIfPlanIsNotOfType(planX.internal, CUDAFTPLANHANDLETYPE::_2dHermetianToReal);
	void* devptrData = (void*)planX.internal.ptrType1;
	auto stride = input.width * sizeof(float) * 2;

	// we need 5 events at most -> host2Device, C2R, normalization, device2Host, finished
	perf.ReserveEvents(5);

	int host2DeviceEvent = perf.InsertCudaEvent();
	cudaError_t cr = cudaMemcpy2D(devptrData, input.width*sizeof(float) * 2, input.ptrData, input.pitch, stride, input.height, cudaMemcpyHostToDevice);
	ThrowIfCudaError(cr, "cudaMemcpy2D");

	int cufftExecC2REvent = perf.InsertCudaEvent();
	cufftResult r = cufftExecC2R(planX.internal.cuFftplan, (cufftComplex*)devptrData, (cufftReal*)devptrData);
	ThrowIfCuFFTError(r, "cufftExecC2R");

	int normalizationEvent = -1;
	if (planX.internal.normalization != 1)
	{
		NppiSize nppiSize = { (int)input.width * 2, (int)input.height };
		normalizationEvent = perf.InsertCudaEvent();
		NppStatus status = nppiDivC_32f_C1IR(planX.internal.normalization, (Npp32f*)devptrData, (int)stride, nppiSize);
		ThrowIfNppError(status, "nppiDivC_32f_C1IR");
	}

	int device2HostEvent = perf.InsertCudaEvent();
	cr = cudaMemcpy2D(output.ptrData, output.pitch, devptrData, stride, output.width*sizeof(float), output.height, cudaMemcpyDeviceToHost);
	ThrowIfCudaError(cr, "cudaMemcpy2D");
	int finishedEvent = perf.InsertCudaEvent();

	perf.SychnronizeToEvent(finishedEvent);

	perf.ReportDurationBetweenEvent(host2DeviceEvent, cufftExecC2REvent, L"Copy Host->Device");
	if (normalizationEvent < 0)
	{
		// w/o normalization
		perf.ReportDurationBetweenEvent(cufftExecC2REvent, device2HostEvent, L"cufftExecC2R");
		perf.ReportDurationBetweenEvent(device2HostEvent, finishedEvent, L"Copy Device->Host");
	}
	else
	{
		// with normalization
		perf.ReportDurationBetweenEvent(cufftExecC2REvent, normalizationEvent, L"cufftExecC2R");
		perf.ReportDurationBetweenEvent(normalizationEvent, device2HostEvent, L"normalize");
		perf.ReportDurationBetweenEvent(device2HostEvent, finishedEvent, L"Copy Device->Host");
	}
}

void cuFft_2d_HermetianToReal(CUDAFTPLANHANDLE plan, const FTLib::FourierTransfComplexInput2d& input, const FTLib::FourierTransfRealOutputData2d& output, FTLib::IPerformanceDataStore* perfData)
{
	if (perfData == nullptr)
	{
		CNullPerfClass npc;
		cuFft_2d_HermetianToReal(plan, input, output, npc);
	}
	else
	{
		CPerfClass pc(0, perfData);
		cuFft_2d_HermetianToReal(plan, input, output, pc);
	}
}

static void DeletePlan_Hermetian2Real(CUDAFTPLANHANDLEDual& planX)
{
	cufftResult r = cufftDestroy(planX.internal.cuFftplan);
	cudaError_t cr = cudaFree((void*)planX.internal.ptrType1);
}
////////////////////////////////////////////////////////////////////////////////

void cuFft_DeletePlan(CUDAFTPLANHANDLE handle)
{
	CUDAFTPLANHANDLEDual& planX = (CUDAFTPLANHANDLEDual&)handle;
	switch (planX.internal.planType)
	{
	case CUDAFTPLANHANDLETYPE::_2dRealToComplexPackedFftw:
		DeletePlan_PackedFftw(planX);
		break;
	case CUDAFTPLANHANDLETYPE::_2dRealToComplex:
		DeletePlan_Complex(planX);
		break;
	case CUDAFTPLANHANDLETYPE::_2dComplexToComplex:
		DeletePlan_Complex2Complex(planX);
		break;
	case CUDAFTPLANHANDLETYPE::_2dHermetianToReal:
		DeletePlan_Hermetian2Real(planX);
		break;
	}
}